use burn::{
    prelude::{Backend, Tensor},
    tensor::Shape,
};
use image::{ImageBuffer, Rgb};
use rand::Rng;
use std::{
    fs::{self, File},
    io::Write,
    path::PathBuf,
};

/// Synthetic dataset for testing YOLO v1
pub struct SyntheticDataset {
    num_samples: usize,
    image_size: (u32, u32),
    num_classes: usize,
    grid_size: usize,
    num_boxes: usize,
}

/// Synthetic data item
pub struct SyntheticItem {
    pub image_path: PathBuf,
    pub targets: Vec<Target>,
}

/// Target representation for YOLO format
#[derive(Debug, Clone)]
pub struct Target {
    pub class_id: usize,
    pub x: f32,      // center x (0-1)
    pub y: f32,      // center y (0-1)
    pub width: f32,  // width (0-1)
    pub height: f32, // height (0-1)
}

impl SyntheticDataset {
    pub fn new(
        num_samples: usize,
        image_size: (u32, u32),
        num_classes: usize,
        grid_size: usize,
        num_boxes: usize,
    ) -> Self {
        Self {
            num_samples,
            image_size,
            num_classes,
            grid_size,
            num_boxes,
        }
    }

    /// Generate synthetic dataset and save to disk
    pub fn generate(&self, output_dir: &str, split: &str) -> Result<Vec<SyntheticItem>, Box<dyn std::error::Error>> {
        let mut rng = rand::thread_rng();
        let mut items = Vec::with_capacity(self.num_samples);

        // Create output directories
        let images_dir = PathBuf::from(output_dir).join("images").join(split);
        let labels_dir = PathBuf::from(output_dir).join("labels").join(split);
        
        fs::create_dir_all(&images_dir)?;
        fs::create_dir_all(&labels_dir)?;

        println!("Generating {} synthetic {} images...", self.num_samples, split);

        for i in 0..self.num_samples {
            // Generate random number of objects (1-3)
            let num_objects = rng.gen_range(1..=3);
            let mut targets = Vec::with_capacity(num_objects);

            // Generate random objects
            for _ in 0..num_objects {
                let class_id = rng.gen_range(0..self.num_classes);
                
                // Generate random bounding box (avoiding edges)
                let width = rng.gen_range(0.1..0.4);
                let height = rng.gen_range(0.1..0.4);
                let x = rng.gen_range(width/2.0..1.0 - width/2.0);
                let y = rng.gen_range(height/2.0..1.0 - height/2.0);

                targets.push(Target {
                    class_id,
                    x,
                    y,
                    width,
                    height,
                });
            }

            // Create synthetic image
            let image = self.generate_image(&targets);
            
            // Save image to file
            let image_filename = format!("synthetic_{}_{}.png", split, i);
            let image_path = images_dir.join(&image_filename);
            image.save(&image_path)?;
            
            // Save YOLO format annotation
            let label_filename = format!("synthetic_{}_{}.txt", split, i);
            let label_path = labels_dir.join(&label_filename);
            self.save_yolo_annotation(&targets, &label_path)?;

            items.push(SyntheticItem {
                image_path,
                targets,
            });
        }

        println!("Generated {} synthetic {} images in {}", self.num_samples, split, output_dir);
        Ok(items)
    }

    /// Generate synthetic image with colored rectangles
    fn generate_image(&self, targets: &[Target]) -> ImageBuffer<Rgb<u8>, Vec<u8>> {
        let (width, height) = self.image_size;
        let mut img = ImageBuffer::from_fn(width, height, |_, _| {
            Rgb([240, 240, 240]) // Light gray background
        });

        let colors = [
            [255, 0, 0],    // Red
            [0, 255, 0],    // Green  
            [0, 0, 255],    // Blue
            [255, 255, 0],  // Yellow
            [255, 0, 255],  // Magenta
        ];

        for (i, target) in targets.iter().enumerate() {
            let color = colors[i % colors.len()];
            
            // Convert normalized coordinates to pixel coordinates
            let center_x = (target.x * width as f32) as u32;
            let center_y = (target.y * height as f32) as u32;
            let bbox_width = (target.width * width as f32) as u32;
            let bbox_height = (target.height * height as f32) as u32;
            
            let start_x = center_x.saturating_sub(bbox_width / 2);
            let start_y = center_y.saturating_sub(bbox_height / 2);
            let end_x = (start_x + bbox_width).min(width - 1);
            let end_y = (start_y + bbox_height).min(height - 1);

            // Draw rectangle
            for x in start_x..end_x {
                for y in start_y..end_y {
                    img.put_pixel(x, y, Rgb(color));
                }
            }
        }

        img
    }

    /// Convert targets to YOLO format tensor
    pub fn targets_to_tensor<B: Backend>(&self, targets: &[Target], device: &B::Device) -> Tensor<B, 3> {
        let shape = Shape::new([self.grid_size, self.grid_size, 5 * self.num_boxes + self.num_classes]);
        let mut data = vec![0.0; shape.num_elements()];

        for target in targets {
            // Calculate grid cell
            let grid_x = (target.x * self.grid_size as f32).floor() as usize;
            let grid_y = (target.y * self.grid_size as f32).floor() as usize;
            
            if grid_x < self.grid_size && grid_y < self.grid_size {
                let cell_index = grid_y * self.grid_size + grid_x;
                let base_index = cell_index * (5 * self.num_boxes + self.num_classes);
                
                // For simplicity, use first bounding box
                if self.num_boxes > 0 {
                    // Bounding box coordinates (relative to cell)
                    let cell_x = target.x * self.grid_size as f32 - grid_x as f32;
                    let cell_y = target.y * self.grid_size as f32 - grid_y as f32;
                    
                    data[base_index] = cell_x;      // x
                    data[base_index + 1] = cell_y;  // y
                    data[base_index + 2] = target.width * self.grid_size as f32;   // width
                    data[base_index + 3] = target.height * self.grid_size as f32;  // height
                    data[base_index + 4] = 1.0;     // confidence
                    
                    // Class probabilities (one-hot encoding)
                    let class_start = base_index + 5 * self.num_boxes;
                    if target.class_id < self.num_classes {
                        data[class_start + target.class_id] = 1.0;
                    }
                }
            }
        }

        Tensor::from_data(data.as_slice(), device)
    }

    /// Save YOLO format annotation file
    fn save_yolo_annotation(&self, targets: &[Target], label_path: &PathBuf) -> Result<(), Box<dyn std::error::Error>> {
        let mut file = File::create(label_path)?;
        
        for target in targets {
            // YOLO format: class_id x_center y_center width height
            let line = format!(
                "{} {:.6} {:.6} {:.6} {:.6}\n",
                target.class_id,
                target.x,
                target.y, 
                target.width,
                target.height
            );
            file.write_all(line.as_bytes())?;
        }
        
        Ok(())
    }

    /// Get dataset size
    pub fn len(&self) -> usize {
        self.num_samples
    }

    /// Check if dataset is empty
    pub fn is_empty(&self) -> bool {
        self.num_samples == 0
    }
}

/// Configuration for synthetic dataset
#[derive(Debug, Clone)]
pub struct SyntheticConfig {
    pub num_train_samples: usize,
    pub num_val_samples: usize,
    pub image_size: (u32, u32),
    pub num_classes: usize,
    pub grid_size: usize,
    pub num_boxes: usize,
}

impl Default for SyntheticConfig {
    fn default() -> Self {
        Self {
            num_train_samples: 100,
            num_val_samples: 20,
            image_size: (448, 448),
            num_classes: 3,  // Simple 3-class dataset for testing
            grid_size: 7,
            num_boxes: 2,
        }
    }
}

/// Create synthetic training dataset
pub fn create_synthetic_train_dataset(config: SyntheticConfig) -> SyntheticDataset {
    SyntheticDataset::new(
        config.num_train_samples,
        config.image_size,
        config.num_classes,
        config.grid_size,
        config.num_boxes,
    )
}

/// Create synthetic validation dataset
pub fn create_synthetic_val_dataset(config: SyntheticConfig) -> SyntheticDataset {
    SyntheticDataset::new(
        config.num_val_samples,
        config.image_size,
        config.num_classes,
        config.grid_size,
        config.num_boxes,
    )
}
